"""
数据预处理：
"""

from typing import Union, Literal
from pathlib import Path
import torchaudio
import numpy as np
from audio_utils import convert_waveform
from data_utils import (
    create_zip,
    cal_gcmvn_stats,
    extract_fbank_features
)
from pyarrow import parquet
import shutil
from io import BytesIO


def preprocess(
        data_dir: Union[Path, str],
        out_dir: Union[Path, str],
        split: Literal["train", "dev", "test"],
        cmvn_type: Literal["global", "local"],
) -> None:

    output_root = Path(out_dir)
    output_root.mkdir(exist_ok=True)

    source_root = output_root / "src"
    source_root.mkdir(exist_ok=True)

    gcmvn_feature_list = []
    for parquet_file in Path(data_dir).glob("*.parquet"):
        pq = parquet.ParquetFile(parquet_file)
        df = pq.read().to_pandas()

        print(parquet_file)
        for idx, row in df.iterrows():
            if idx % 10000 == 0:
                print(idx)
            audio_io = BytesIO(row['src_audio']['bytes'])
            waveform, sample_rate = torchaudio.load(audio_io)
            src_sample_rate = 16_000
            waveform, sample_rate = convert_waveform(
                waveform, sample_rate, to_mono=True,
                to_sample_rate=src_sample_rate
            )
            utt_id = str(row['src_audio']['path']).split('.')[0]
            features = extract_fbank_features(
                waveform, sample_rate, output_path=source_root / f"{utt_id}.npy"
            )
            if split == 'train' and cmvn_type == "global":
                gcmvn_feature_list.append(features)

    if split == 'train' and cmvn_type == "global":
        stats = cal_gcmvn_stats(gcmvn_feature_list)
        with open(output_root / "gcmvn.npz", "wb") as fw:
            np.savez(fw, mean=stats["mean"], std=stats["std"])

    print("ZIPing source audios/features...")
    source_zip_path = output_root / f"{source_root.name}.zip"
    create_zip(source_root, source_zip_path)
    shutil.rmtree(source_root)


if __name__ == '__main__':
    data_dir = r'/root/autodl-tmp/malay/train'
    out_dir = r'/root/autodl-tmp/malay/train_processed'
    preprocess(data_dir, out_dir, split="train", cmvn_type='global')

